%% @doc RabbitMQ 生产者模块。
-module(emqx_plugin_rabbitmq_producer).

-behaviour(emqx_resource).

-include_lib("emqx_resource/include/emqx_resource.hrl").
-include_lib("emqx/include/logger.hrl").
-include("emqx_plugin_rabbitmq.hrl").

%% amqp_client records
-include_lib("amqp_client/include/amqp_client.hrl").

-export([
    query_mode/1,
    callback_mode/0,
    on_start/2,
    on_get_status/2,
    on_stop/2,
    on_query_async/4
]).

-export([schedule_reconnect/2]).

-define(MAX_RETRIES, 5).
-define(RETRY_INTERVAL, 5000). % 5 seconds

%% @doc 查询模式 - 使用简单异步内部缓冲。
query_mode(_) ->
    simple_async_internal_buffer.

%% @doc 回调模式 - 尽可能使用异步。
callback_mode() ->
    async_if_possible.

%% @doc 启动回调函数。
on_start(InstId, Config) ->
    ?SLOG(info, #{msg => "rabbitmq_producer_on_start", inst_id => InstId, config => Config}),
    ConnectionParams = maps:get(connection, Config, #{}),

    HostsStr = maps:get(hosts, ConnectionParams, "localhost:5672"),
    HostsList = parse_hosts(HostsStr),

    Username = maps:get(username, ConnectionParams, <<"guest">>),
    Password = emqx_secret:unwrap(maps:get(password, ConnectionParams, <<"guest">>)),
    VirtualHost = maps:get(virtual_host, ConnectionParams, <<"/">>),
    Heartbeat = maps:get(heartbeat, ConnectionParams, 60),

    State = #{
        inst_id => InstId,
        config => Config,
        host => case HostsList of [{H,_}|_] -> H; [] -> "localhost" end,
        port => case HostsList of [{_,P}|_] -> P; [] -> 5672 end,
        username => Username,
        password => Password,
        virtual_host => VirtualHost,
        heartbeat => Heartbeat,
        conn => undefined,
        channel => undefined,
        retry_count => 0,
        connection_status => disconnected
    },

    case connect_and_open_channel(State) of
        {ok, NewState} ->
            {ok, NewState#{connection_status => connected, retry_count => 0}};
        {error, Reason, NewState} ->
            ?SLOG(error, #{msg => "rabbitmq_producer_initial_connection_failed", inst_id => InstId, reason => Reason}),
            schedule_reconnect(InstId, NewState),
            {ok, NewState#{connection_status => disconnected, retry_count => 1}}
    end.

%% @doc 获取状态回调函数。
on_get_status(InstId, #{conn := Conn, channel := Channel, retry_count := RetryCount} = _State) ->
    IsConnOpen = is_pid(Conn) andalso erlang:is_process_alive(Conn),
    IsChanOpen = is_pid(Channel) andalso erlang:is_process_alive(Channel),
    if
        RetryCount >= ?MAX_RETRIES andalso RetryCount > 0 ->
            ?SLOG(error, #{msg => "rabbitmq_producer_max_retries_reached", inst_id => InstId, retry_count => RetryCount}),
            ?status_disconnected;
        IsConnOpen andalso IsChanOpen ->
            ?status_connected;
        true ->
            ?status_connecting
    end;
on_get_status(InstId, State) ->
    ?SLOG(debug, #{msg => "rabbitmq_producer_on_get_status_fallback", inst_id => InstId, state_keys => maps:keys(State)}),
    ?status_connecting.


%% @doc 停止回调函数。
on_stop(InstId, #{conn := Conn, channel := Channel} = State) ->
    ?SLOG(info, #{msg => "rabbitmq_producer_on_stop", inst_id => InstId}),
    spawn(fun() ->
        close_channel_and_connection(Channel, Conn, InstId)
    end),
    persistent_term:erase({?EMQX_PLUGIN_RABBITMQ_APP, ?EMQX_PLUGIN_RABBITMQ_PRODUCERS}),
    {ok, State}.

%% @doc 异步查询回调函数 (emqx_resource behavior)。
on_query_async(InstId, _Query = {publish, Exchange, RoutingKey, Persistent, Payload, Context}, _From, State = #{channel := _OrigChannel, config := _OrigConfig}) when is_map(Context) ->
    ?SLOG(info, #{msg => "rabbitmq_producer_on_query_async", inst_id => InstId, exchange => Exchange, routing_key => RoutingKey, context => Context}),
    
    NewState = ensure_channel_is_ready(InstId, State), % 确保通道准备就绪

    case maps:get(channel, NewState, undefined) of
        ReadyChannel when is_pid(ReadyChannel) -> % 首先确保是PID
            case erlang:is_process_alive(ReadyChannel) of % 然后检查是否存活
                true ->
                    DeliveryMode = case Persistent of
                                       true -> 2;
                                       false -> 1
                                   end,
                    AppId = atom_to_binary(?EMQX_PLUGIN_RABBITMQ_APP, utf8),
                    ContentType = <<"application/octet-stream">>,
                    Timestamp = erlang:system_time(second),
                    AMQPMessage = #amqp_msg{
                        payload = Payload,
                        props = #'P_basic'{
                            app_id = AppId,
                            content_type = ContentType,
                            delivery_mode = DeliveryMode,
                            timestamp = Timestamp
                        }
                    },
                    try amqp_channel:cast(ReadyChannel, #'basic.publish'{exchange = Exchange, routing_key = RoutingKey}, AMQPMessage) of
                        _ ->
                            ?SLOG(info, #{msg => "rabbitmq_publish_successful", inst_id => InstId, exchange => Exchange, routing_key => RoutingKey}),
                            %% 修改返回格式，使用正确的格式
                            ok
                    catch
                        Error:Reason:Stacktrace ->
                            ?SLOG(error, #{msg => "rabbitmq_publish_failed", inst_id => InstId, exchange => Exchange, routing_key => RoutingKey, error => Error, reason => Reason, stacktrace => Stacktrace}),
                            ClosedState = close_and_schedule_reconnect(InstId, NewState),
                            {error, {publish_failed, Reason}, ClosedState}
                    end;
                false -> % PID 但不存活
                    ?SLOG(error, #{msg => "rabbitmq_channel_not_alive_for_publish", inst_id => InstId, exchange => Exchange, routing_key => RoutingKey}),
                    ReconnectingState = schedule_reconnect_if_needed(InstId, NewState),
                    {error, channel_not_ready, ReconnectingState}
            end;
        InvalidChannel -> % 不是 PID (可能为 undefined 或其他)
            ?SLOG(error, #{msg => "rabbitmq_channel_invalid_for_publish", inst_id => InstId, exchange => Exchange, routing_key => RoutingKey, channel_val => InvalidChannel}),
            ReconnectingState = schedule_reconnect_if_needed(InstId, NewState),
            {error, channel_not_ready, ReconnectingState}
    end;

%% 添加一个新的子句，处理{Querys, Context}格式的查询
on_query_async(InstId, {Querys, Context}, _From, State) when is_list(Querys), is_map(Context) ->
    ?SLOG(info, #{msg => "rabbitmq_producer_on_query_async_list", inst_id => InstId, querys_count => length(Querys)}),
    %% 从配置中获取默认值
    #{config := Config} = State,
    ProducerConfig = maps:get(producer, Config, #{}),
    DefaultExchange = maps:get(default_exchange, ProducerConfig, <<>>),
    DefaultRoutingKey = maps:get(default_routing_key, ProducerConfig, <<"emqx_message">>),
    PersistentMessages = maps:get(persistent_messages, ProducerConfig, true),
    
    %% 处理每个查询
    Results = lists:map(
        fun({MessageProps, Payload}) when is_map(MessageProps) ->
            Exchange = maps:get(exchange, MessageProps, DefaultExchange),
            RoutingKey = maps:get(routing_key, MessageProps, DefaultRoutingKey),
            %% 直接处理单个消息，而不是递归调用
            publish_single_message(InstId, Exchange, RoutingKey, PersistentMessages, Payload, Context, State);
        ({Topic, Payload}) when is_binary(Topic) ->
            %% 简单格式，使用默认交换机和主题作为路由键
            publish_single_message(InstId, DefaultExchange, Topic, PersistentMessages, Payload, Context, State);
        (UnknownFormat) ->
            ?SLOG(error, #{msg => "rabbitmq_unknown_query_format", inst_id => InstId, query => UnknownFormat}),
            {error, unknown_query_format}
        end,
        Querys
    ),
    
    %% 检查结果，如果有任何错误，返回第一个错误
    case lists:keyfind(error, 1, Results) of
        {error, Reason} ->
            {error, Reason, State};
        false ->
            %% 修改返回格式，使用正确的格式
            ok
    end;

%% 保留原有的函数子句作为备用
on_query_async(InstId, UnknownQuery, _From, State) ->
    ?SLOG(warning, #{msg => "rabbitmq_producer_unknown_query", inst_id => InstId, query => UnknownQuery}),
    {error, unknown_query, State}.

%% 新增辅助函数，处理单个消息发布
publish_single_message(InstId, Exchange, RoutingKey, Persistent, Payload, Context, State = #{channel := Channel}) ->
    case is_pid(Channel) andalso erlang:is_process_alive(Channel) of
        true ->
            DeliveryMode = case Persistent of
                               true -> 2;
                               false -> 1
                           end,
            AppId = atom_to_binary(?EMQX_PLUGIN_RABBITMQ_APP, utf8),
            ContentType = <<"application/octet-stream">>,
            Timestamp = erlang:system_time(second),
            AMQPMessage = #amqp_msg{
                payload = Payload,
                props = #'P_basic'{
                    app_id = AppId,
                    content_type = ContentType,
                    delivery_mode = DeliveryMode,
                    timestamp = Timestamp
                }
            },
            try amqp_channel:cast(Channel, #'basic.publish'{exchange = Exchange, routing_key = RoutingKey}, AMQPMessage) of
                _ ->
                    ?SLOG(debug, #{msg => "rabbitmq_single_publish_successful", inst_id => InstId, exchange => Exchange, routing_key => RoutingKey}),
                    ok
            catch
                Error:Reason:Stacktrace ->
                    ?SLOG(error, #{msg => "rabbitmq_single_publish_failed", inst_id => InstId, exchange => Exchange, routing_key => RoutingKey, error => Error, reason => Reason, stacktrace => Stacktrace}),
                    {error, {publish_failed, Reason}}
            end;
        false ->
            ?SLOG(error, #{msg => "rabbitmq_channel_not_ready_for_single_publish", inst_id => InstId, exchange => Exchange, routing_key => RoutingKey}),
            {error, channel_not_ready}
    end.

% 新增辅助函数，确保通道准备就绪
ensure_channel_is_ready(InstId, State = #{channel := Channel, conn := Conn}) ->
    IsConnAlive = is_pid(Conn) andalso erlang:is_process_alive(Conn),
    IsChanAlive = is_pid(Channel) andalso erlang:is_process_alive(Channel),

    if
        IsConnAlive andalso IsChanAlive ->
            State; % 通道和连接都正常
        IsConnAlive andalso not IsChanAlive -> % 连接正常，通道异常
            ?SLOG(warning, #{msg => "rabbitmq_channel_dead_but_connection_alive_closing_both_and_reconnecting", inst_id => InstId}),
            close_channel_silently(Channel, InstId), % 先尝试安静地关闭旧通道句柄
            close_connection_silently(Conn, InstId), % 也关闭连接
            % 像连接也断开一样处理，进行完整的重连
            case connect_and_open_channel(State) of
                {ok, NewStateConnected} ->
                    NewStateConnected#{connection_status => connected, retry_count => 0};
                {error, Reason, NewStateDisconnected} ->
                    ?SLOG(error, #{msg => "rabbitmq_reconnect_attempt_failed_after_channel_dead", inst_id => InstId, reason => Reason}),
                    schedule_reconnect_if_needed(InstId, NewStateDisconnected#{channel := undefined, conn:= undefined, connection_status => disconnected})
            end;
        not IsConnAlive -> % 连接也断了
            ?SLOG(warning, #{msg => "rabbitmq_connection_dead_attempting_reconnect", inst_id => InstId}),
            close_channel_silently(Channel, InstId), % 清理旧通道
            close_connection_silently(Conn, InstId), % 清理旧连接
            case connect_and_open_channel(State) of
                {ok, NewStateConnected} ->
                    NewStateConnected#{connection_status => connected, retry_count => 0};
                {error, Reason, NewStateDisconnected} ->
                    ?SLOG(error, #{msg => "rabbitmq_reconnect_attempt_failed_in_ensure_channel", inst_id => InstId, reason => Reason}),
                    schedule_reconnect_if_needed(InstId, NewStateDisconnected#{channel := undefined, conn:= undefined, connection_status => disconnected})
            end
    end.

close_and_schedule_reconnect(InstId, State = #{channel := Channel, conn := Conn}) ->
    close_channel_silently(Channel, InstId),
    close_connection_silently(Conn, InstId),
    schedule_reconnect_if_needed(InstId, State#{channel := undefined, conn := undefined, connection_status => disconnected}).

schedule_reconnect_if_needed(InstId, State = #{retry_count := RetryCount, connection_status := disconnected}) ->
    % 只有在当前状态是 disconnected 且重试次数未达上限时才调度
    if RetryCount < ?MAX_RETRIES ->
        schedule_reconnect(InstId, State);
        true -> State % 已达上限或状态不是 disconnected，不再调度
    end;
schedule_reconnect_if_needed(_InstId, State) -> % 其他状态，不调度
    State.

close_channel_silently(Channel, InstId) ->
    try amqp_channel:close(Channel) catch _:_ -> ok end,
    ?SLOG(debug, #{msg => "rabbitmq_channel_closed_silently", inst_id => InstId, channel => Channel}).

close_connection_silently(Conn, InstId) ->
    try amqp_connection:close(Conn) catch _:_ -> ok end,
    ?SLOG(debug, #{msg => "rabbitmq_connection_closed_silently", inst_id => InstId, conn => Conn}).


%% Internal connect and open channel
connect_and_open_channel(#{host := Host, port := Port, username := UsernameIn,
                          password := PasswordIn, virtual_host := VHostIn,
                          heartbeat := Heartbeat} = State) ->
    UsernameBin = case is_binary(UsernameIn) of
                      true -> UsernameIn;
                      false -> list_to_binary(UsernameIn)
                  end,
    PasswordBin = case is_binary(PasswordIn) of
                      true -> PasswordIn;
                      false -> list_to_binary(PasswordIn)
                  end,
    VHostBin = case is_binary(VHostIn) of
                   true -> VHostIn;
                   false -> list_to_binary(VHostIn)
               end,
    Params = #amqp_params_network{
        host = Host,
        port = Port,
        username = UsernameBin,
        password = PasswordBin,
        virtual_host = VHostBin,
        heartbeat = Heartbeat
    },
    try
        case amqp_connection:start(Params) of
            {ok, Conn} ->
                case amqp_connection:open_channel(Conn) of
                    {ok, Channel} ->
                        {ok, State#{conn => Conn, channel => Channel}};
                    {error, ChanReason} ->
                        {error, ChanReason, State}
                end;
            {error, ConnReason} ->
                {error, ConnReason, State}
        end
    catch
        Error:Reason:Stack ->
            ?SLOG(error, #{msg => "rabbitmq_connect_exception",
                          error => Error,
                          reason => Reason,
                          stacktrace => Stack,
                          function => ?FUNCTION_NAME}),
            {error, Reason, State}
    end.

%% @doc 关闭通道和连接。
close_channel_and_connection(Channel, Conn, InstId) ->
    case is_pid(Channel) of
        true ->
            try amqp_channel:close(Channel) of
                ok -> ?SLOG(info, #{msg => "rabbitmq_channel_closed", inst_id => InstId});
                {error, CloseChanReason} -> ?SLOG(error, #{msg => "rabbitmq_channel_close_failed", inst_id => InstId, reason => CloseChanReason})
            catch
                Error1:Reason1:Stack1 ->
                    ?SLOG(error, #{msg => "rabbitmq_channel_close_exception", inst_id => InstId, error=>Error1, reason => Reason1, stacktrace => Stack1})
            end;
        false -> ok
    end,
    case is_pid(Conn) of
        true ->
            try amqp_connection:close(Conn) of
                ok -> ?SLOG(info, #{msg => "rabbitmq_connection_closed", inst_id => InstId});
                {error, CloseConnReason} -> ?SLOG(error, #{msg => "rabbitmq_connection_close_failed", inst_id => InstId, reason => CloseConnReason})
            catch
                Error2:Reason2:Stack2 ->
                    ?SLOG(error, #{msg => "rabbitmq_connection_close_exception", inst_id => InstId, error=>Error2, reason => Reason2, stacktrace => Stack2})
            end;
        false -> ok
    end.


%% @doc 调度重连任务。
schedule_reconnect(InstId, State = #{retry_count := RetryCount}) ->
    TargetPid = case erlang:whereis(emqx_plugin_rabbitmq_sup) of
        undefined ->
            ?SLOG(warning, #{msg => "emqx_plugin_rabbitmq_sup_not_found_using_self_for_timer", inst_id => InstId}),
            self();
        Pid -> Pid
    end,
    if
        RetryCount < ?MAX_RETRIES ->
            ?SLOG(info, #{msg => "rabbitmq_producer_schedule_reconnect", inst_id => InstId, retry_in_ms => ?RETRY_INTERVAL, retry_count => RetryCount + 1}),
            erlang:send_after(?RETRY_INTERVAL, TargetPid, {reconnect, InstId, State#{retry_count := RetryCount + 1}});
        true ->
            ?SLOG(error, #{msg => "rabbitmq_producer_max_retries_reached_cannot_reconnect", inst_id => InstId, retry_count => RetryCount})
    end.

%% 解析 hosts 字符串为 {Host, Port} 元组列表
parse_hosts(HostsStr) when is_binary(HostsStr) ->
    parse_hosts(binary_to_list(HostsStr));
parse_hosts(HostsStr) when is_list(HostsStr) ->
    Hosts = string:tokens(HostsStr, ","),
    [parse_host(Host) || Host <- Hosts].

%% 解析单个 host:port 字符串
parse_host(HostPort) ->
    case string:tokens(HostPort, ":") of
        [Host, PortStr] ->
            {Host, list_to_integer(PortStr)};
        [Host] ->
            {Host, 5672}  % 默认 RabbitMQ 端口
    end.

% 新增声明交换机的辅助函数
declare_exchange(Channel, Exchange, Type) ->
    ?SLOG(info, #{msg => "exchange_declare_skipped_manual_creation", exchange => Exchange, type => Type, channel => Channel}),
    ok.

%% 新增辅助函数：将二进制类型的交换机类型转换为原子类型
translate_exchange_type(<<"direct">>) -> 'direct';
translate_exchange_type(<<"topic">>) -> 'topic';
translate_exchange_type(<<"fanout">>) -> 'fanout';
translate_exchange_type(<<"headers">>) -> 'headers';
translate_exchange_type(_) -> 'fanout'. %% 默认使用fanout类型
